﻿using FlyRemote.Base;
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text;
using static FlyRemote.Remote.Native.Enumerations;
using static FlyRemote.Remote.Native.Structures;

namespace FlyRemote.Remote
{
    public abstract class MemoryAssistBase : IMemoryAssist
    {
        public abstract string Description();


        /// <summary>
        /// The chunk size for memory regions. Prevents large allocations.
        /// </summary>
        protected const int ChunkSize = 2000000000;


        /// <summary>
        /// Gets or sets the module cache of process modules.
        /// </summary>
        protected TtlCache<Int32, List<NormalizedModule>> ModuleCache { get; set; }



        /// <summary>
        /// Gets a reference to the target process. This is an optimization to minimize accesses to the Processes component of the Engine.
        /// </summary>
        public Process ExternalProcess { get; set; }

        public void Update(Process process)
        {
            this.ExternalProcess = process;
        }






        #region 内存查询

        public ulong AddressToModule(ulong address, out string moduleName)
        {
            NormalizedModule containingModule = GetModules()
                .Select(module => module)
                .Where(module => module.ContainsAddress(address))
                .FirstOrDefault();

            moduleName = containingModule?.Name ?? String.Empty;

            return containingModule == null ? address : address - containingModule.BaseAddress;

        }



        public ulong EvaluatePointer(ulong address, IEnumerable<int> offsets)
        {
            throw new NotImplementedException();
        }


        public IEnumerable<NormalizedRegion> GetAllVirtualPages()
        {
            throw new NotImplementedException();
        }

        public IEnumerable<NormalizedRegion> GetHeapAddresses()
        {
            throw new NotImplementedException();
        }

        public ulong GetMaximumAddress()
        {
            throw new NotImplementedException();
        }

        public ulong GetMaxUsermodeAddress()
        {
            if (PAFactory.Default.IsOpenedProcess32Bit())
            {
                // 即 7FFF FFFF
                return Int32.MaxValue;
            }
            else
            {
                // 在本机上是 7ff512340000 以后是模块空间
                // 即前面空 4个0 
                return 0x00007FFF00000000;
            }
        }

        public ulong GetMinUsermodeAddress()
        {
            return UInt16.MaxValue;
        }



        /// <summary>
        /// 此处已修改为  使用 peb 来获取  
        /// </summary>
        public IEnumerable<NormalizedModule> GetModules()
        {
            return PebHelper.GetModules(this);
        }

        /// <summary>
        ///  此处已修改为  使用 peb 来获取
        /// </summary>
        public IEnumerable<NormalizedRegion> GetStackAddresses()
        {
            return PebHelper.GetStacks(this);
        }




        private static IEnumerable<MemoryBasicInformation64> VirtualPages(
            IntPtr processHandle,
            UInt64 startAddress,
            UInt64 endAddress,
            MemoryProtectionFlags requiredProtection,
            MemoryProtectionFlags excludedProtection,
            MemoryTypeEnum allowedTypes)
        {
            if (startAddress >= endAddress)
            {
                yield return new MemoryBasicInformation64();
            }

            Boolean wrappedAround = false;
            int queryResult;

            // Enumerate the memory pages
            do
            {
                // Allocate the structure to store information of memory
                MemoryBasicInformation64 memoryInfo = new MemoryBasicInformation64();

                if (!Environment.Is64BitProcess)
                {
                    // 32 Bit struct is not the same
                    MemoryBasicInformation32 memoryInfo32 = new MemoryBasicInformation32();

                    // Query the memory region (32 bit native method)
                    queryResult = Native.NativeMethods.VirtualQueryEx(processHandle, startAddress.ToIntPtr(), out memoryInfo32, Marshal.SizeOf(memoryInfo32));

                    // Copy from the 32 bit struct to the 64 bit struct
                    memoryInfo.AllocationBase = memoryInfo32.AllocationBase;
                    memoryInfo.AllocationProtect = memoryInfo32.AllocationProtect;
                    memoryInfo.BaseAddress = memoryInfo32.BaseAddress;
                    memoryInfo.Protect = memoryInfo32.Protect;
                    memoryInfo.RegionSize = memoryInfo32.RegionSize;
                    memoryInfo.State = memoryInfo32.State;
                    memoryInfo.Type = memoryInfo32.Type;
                }
                else
                {
                    // Query the memory region (64 bit native method)
                    queryResult = Native.NativeMethods.VirtualQueryEx(processHandle, startAddress.ToIntPtr(), out memoryInfo, Marshal.SizeOf(memoryInfo));
                }

                // Increment the starting address with the size of the page
                UInt64 previousFrom = startAddress;
                startAddress = startAddress.Add(memoryInfo.RegionSize);

                if (previousFrom > startAddress)
                {
                    wrappedAround = true;
                }

                // Ignore free memory. These are unallocated memory regions.
                if ((memoryInfo.State & MemoryStateFlags.Free) != 0)
                {
                    continue;
                }

                // At least one readable memory flag is required
                if ((memoryInfo.Protect & MemoryProtectionFlags.ReadOnly) == 0 && (memoryInfo.Protect & MemoryProtectionFlags.ExecuteRead) == 0 &&
                    (memoryInfo.Protect & MemoryProtectionFlags.ExecuteReadWrite) == 0 && (memoryInfo.Protect & MemoryProtectionFlags.ReadWrite) == 0)
                {
                    continue;
                }

                // Do not bother with this shit, this memory is not worth scanning
                if ((memoryInfo.Protect & MemoryProtectionFlags.ZeroAccess) != 0 || (memoryInfo.Protect & MemoryProtectionFlags.NoAccess) != 0 || (memoryInfo.Protect & MemoryProtectionFlags.Guard) != 0)
                {
                    continue;
                }

                // Enforce allowed types
                switch (memoryInfo.Type)
                {
                    case MemoryTypeFlags.None:
                        if ((allowedTypes & MemoryTypeEnum.None) == 0)
                        {
                            continue;
                        }

                        break;
                    case MemoryTypeFlags.Private:
                        if ((allowedTypes & MemoryTypeEnum.Private) == 0)
                        {
                            continue;
                        }

                        break;
                    case MemoryTypeFlags.Image:
                        if ((allowedTypes & MemoryTypeEnum.Image) == 0)
                        {
                            continue;
                        }

                        break;
                    case MemoryTypeFlags.Mapped:
                        if ((allowedTypes & MemoryTypeEnum.Mapped) == 0)
                        {
                            continue;
                        }

                        break;
                }

                // Ensure at least one required protection flag is set
                if (requiredProtection != 0 && (memoryInfo.Protect & requiredProtection) == 0)
                {
                    continue;
                }

                // Ensure no ignored protection flags are set
                if (excludedProtection != 0 && (memoryInfo.Protect & excludedProtection) != 0)
                {
                    continue;
                }

                // Return the memory page
                yield return memoryInfo;
            }
            while (startAddress < endAddress && queryResult != 0 && !wrappedAround);
        }

        public IList<NormalizedRegion> GetVirtualPages(MemoryProtectionEnum requiredProtection, MemoryProtectionEnum excludedProtection, MemoryTypeEnum allowedTypes, ulong startAddress, ulong endAddress)
        {
            MemoryProtectionFlags requiredFlags = 0;
            MemoryProtectionFlags excludedFlags = 0;

            if ((requiredProtection & MemoryProtectionEnum.Write) != 0)
            {
                requiredFlags |= MemoryProtectionFlags.ExecuteReadWrite;
                requiredFlags |= MemoryProtectionFlags.ReadWrite;
            }

            if ((requiredProtection & MemoryProtectionEnum.Execute) != 0)
            {
                requiredFlags |= MemoryProtectionFlags.Execute;
                requiredFlags |= MemoryProtectionFlags.ExecuteRead;
                requiredFlags |= MemoryProtectionFlags.ExecuteReadWrite;
                requiredFlags |= MemoryProtectionFlags.ExecuteWriteCopy;
            }

            if ((requiredProtection & MemoryProtectionEnum.CopyOnWrite) != 0)
            {
                requiredFlags |= MemoryProtectionFlags.WriteCopy;
                requiredFlags |= MemoryProtectionFlags.ExecuteWriteCopy;
            }

            if ((excludedProtection & MemoryProtectionEnum.Write) != 0)
            {
                excludedFlags |= MemoryProtectionFlags.ExecuteReadWrite;
                excludedFlags |= MemoryProtectionFlags.ReadWrite;
            }

            if ((excludedProtection & MemoryProtectionEnum.Execute) != 0)
            {
                excludedFlags |= MemoryProtectionFlags.Execute;
                excludedFlags |= MemoryProtectionFlags.ExecuteRead;
                excludedFlags |= MemoryProtectionFlags.ExecuteReadWrite;
                excludedFlags |= MemoryProtectionFlags.ExecuteWriteCopy;
            }

            if ((excludedProtection & MemoryProtectionEnum.CopyOnWrite) != 0)
            {
                excludedFlags |= MemoryProtectionFlags.WriteCopy;
                excludedFlags |= MemoryProtectionFlags.ExecuteWriteCopy;
            }

            IEnumerable<MemoryBasicInformation64> memoryInfo = VirtualPages(this.ExternalProcess == null ? IntPtr.Zero : this.ExternalProcess.Handle, startAddress, endAddress, requiredFlags, excludedFlags, allowedTypes);

            IList<NormalizedRegion> regions = new List<NormalizedRegion>(0x500);

            foreach (MemoryBasicInformation64 next in memoryInfo)
            {

                if (next.RegionSize < ChunkSize)
                {
                    regions.Add(new NormalizedRegion(next.BaseAddress.ToUInt64(), next.RegionSize.ToInt32()));
                }
                else
                {
                    // This region requires chunking
                    Int64 remaining = next.RegionSize;
                    UInt64 currentBaseAddress = next.BaseAddress.ToUInt64();

                    while (remaining >= ChunkSize)
                    {
                        regions.Add(new NormalizedRegion(currentBaseAddress, ChunkSize));

                        remaining -= ChunkSize;
                        currentBaseAddress = currentBaseAddress.Add(ChunkSize, wrapAround: false);
                    }

                    if (remaining > 0)
                    {
                        regions.Add(new NormalizedRegion(currentBaseAddress, remaining.ToInt32()));
                    }
                }
            }

            return regions;
        }

        public void EnumVirtualPages(MemoryProtectionEnum requiredProtection, MemoryProtectionEnum excludedProtection, MemoryTypeEnum allowedTypes, ulong startAddress, ulong endAddress, VirtualPageEnumProc proc)
        {

            MemoryProtectionFlags requiredFlags = 0;
            MemoryProtectionFlags excludedFlags = 0;

            if ((requiredProtection & MemoryProtectionEnum.Write) != 0)
            {
                requiredFlags |= MemoryProtectionFlags.ExecuteReadWrite;
                requiredFlags |= MemoryProtectionFlags.ReadWrite;
            }

            if ((requiredProtection & MemoryProtectionEnum.Execute) != 0)
            {
                requiredFlags |= MemoryProtectionFlags.Execute;
                requiredFlags |= MemoryProtectionFlags.ExecuteRead;
                requiredFlags |= MemoryProtectionFlags.ExecuteReadWrite;
                requiredFlags |= MemoryProtectionFlags.ExecuteWriteCopy;
            }

            if ((requiredProtection & MemoryProtectionEnum.CopyOnWrite) != 0)
            {
                requiredFlags |= MemoryProtectionFlags.WriteCopy;
                requiredFlags |= MemoryProtectionFlags.ExecuteWriteCopy;
            }

            if ((excludedProtection & MemoryProtectionEnum.Write) != 0)
            {
                excludedFlags |= MemoryProtectionFlags.ExecuteReadWrite;
                excludedFlags |= MemoryProtectionFlags.ReadWrite;
            }

            if ((excludedProtection & MemoryProtectionEnum.Execute) != 0)
            {
                excludedFlags |= MemoryProtectionFlags.Execute;
                excludedFlags |= MemoryProtectionFlags.ExecuteRead;
                excludedFlags |= MemoryProtectionFlags.ExecuteReadWrite;
                excludedFlags |= MemoryProtectionFlags.ExecuteWriteCopy;
            }

            if ((excludedProtection & MemoryProtectionEnum.CopyOnWrite) != 0)
            {
                excludedFlags |= MemoryProtectionFlags.WriteCopy;
                excludedFlags |= MemoryProtectionFlags.ExecuteWriteCopy;
            }

            IEnumerable<MemoryBasicInformation64> memoryInfo = VirtualPages(this.ExternalProcess == null ? IntPtr.Zero : this.ExternalProcess.Handle, startAddress, endAddress, requiredFlags, excludedFlags, allowedTypes);

            foreach (MemoryBasicInformation64 next in memoryInfo)
            {
                unchecked
                {
                    proc((ulong)next.BaseAddress, (uint)next.RegionSize);
                }
            }
        }




        public ulong ResolveModule(string identifier)
        {
            UInt64 result = 0;

            identifier = identifier?.RemoveSuffixes(true, ".exe", ".dll");
            IEnumerable<NormalizedModule> modules = this.GetModules()
                ?.ToList()
                ?.Select(module => module)
                ?.Where(module => module.Name.RemoveSuffixes(true, ".exe", ".dll").Equals(identifier, StringComparison.OrdinalIgnoreCase));

            if (modules.Count() > 0)
            {
                result = modules.First().BaseAddress;
            }

            return result;
        }

        #endregion

        #region 内存分配
        public abstract ulong AllocateMemory(int size);

        public abstract ulong AllocateMemory(int size, ulong allocAddress);

        public abstract void FreeMemory(ulong address);



        #endregion


        #region 内存读写 
        public object Read(DataType dataType, ulong address, out bool success)
        {

            if (address == 0)
            {
                success = false;
                return null;
            }

            Object value = null;

            switch (dataType.Code)
            {
                case DataType.TypeCode.u8:
                    value = this.Read<Byte>(address, out success);
                    break;
                case DataType.TypeCode.i8:
                    value = this.Read<SByte>(address, out success);
                    break;
                case DataType.TypeCode.i16:
                    value = this.Read<Int16>(address, out success);
                    break;
                case DataType.TypeCode.u16:
                    value = this.Read<UInt16>(address, out success);
                    break;
                case DataType.TypeCode.i32:
                    value = this.Read<Int32>(address, out success);
                    break;
                case DataType.TypeCode.u32:
                    value = this.Read<UInt32>(address, out success);
                    break;
                case DataType.TypeCode.i64:
                    value = this.Read<Int64>(address, out success);
                    break;
                case DataType.TypeCode.u64:
                    value = this.Read<UInt64>(address, out success);
                    break;
                case DataType.TypeCode.f32:
                    value = this.Read<Single>(address, out success);
                    break;
                case DataType.TypeCode.f64:
                    value = this.Read<Double>(address, out success);
                    break;
                case DataType.TypeCode.Ansi:
                    {
                        var arr = this.ReadBytes(address, dataType.Size, out success);
                        if (success)
                            value = Encoding.ASCII.GetString(arr);
                        break;
                    }
                case DataType.TypeCode.Utf8:
                    {
                        var arr = this.ReadBytes(address, dataType.Size, out success);
                        if (success)
                            value = Encoding.UTF8.GetString(arr);
                        break;
                    }
                case DataType.TypeCode.Utf16:
                    {
                        var arr = this.ReadBytes(address, dataType.Size, out success);
                        if (success)
                            value = Encoding.Unicode.GetString(arr);
                        break;
                    }

                case DataType.TypeCode.ByteArray:
                    {
                        var arr = this.ReadBytes(address, dataType.Size, out success);
                        value = arr;
                        break;
                    }



                default:
                    value = "?";
                    success = false;
                    break;
            }

            if (!success)
            {
                value = "?";
            }

            return value;
        }

        public T Read<T>(ulong address, out bool success)
        {
            Byte[] byteArray = this.ReadBytes(address, Marshal.SizeOf(typeof(T)), out success);

            if (byteArray == null)
                return default(T);

            return DataType.FromBytes<T>(byteArray);
        }


        public void Write(DataType elementType, ulong address, object value)
        {
            Byte[] bytes;

            switch (elementType)
            {
                case DataType type when type == DataType.u8:
                    bytes = new Byte[] { (Byte)value };
                    break;
                case DataType type when type == DataType.i8:
                    bytes = new Byte[] { unchecked((Byte)(SByte)value) };
                    break;
                case DataType type when type == DataType.i16:
                    bytes = BitConverter.GetBytes((Int16)value);
                    break;
                case DataType type when type == DataType.i32:
                    bytes = BitConverter.GetBytes((Int32)value);
                    break;
                case DataType type when type == DataType.i64:
                    bytes = BitConverter.GetBytes((Int64)value);
                    break;
                case DataType type when type == DataType.u16:
                    bytes = BitConverter.GetBytes((UInt16)value);
                    break;
                case DataType type when type == DataType.u32:
                    bytes = BitConverter.GetBytes((UInt32)value);
                    break;
                case DataType type when type == DataType.u64:
                    bytes = BitConverter.GetBytes((UInt64)value);
                    break;
                case DataType type when type == DataType.f32:
                    bytes = BitConverter.GetBytes((Single)value);
                    break;
                case DataType type when type == DataType.f64:
                    bytes = BitConverter.GetBytes((Double)value);
                    break;
                default:
                    throw new ArgumentException("Invalid type provided");
            }

            this.WriteBytes(address, bytes);

        }

        public void Write<T>(ulong address, T value)
        {

            this.Write(typeof(T), address, (Object)value);

        }
        public abstract byte[] ReadBytes(ulong address, int size, out bool success);

        public abstract bool WriteBytes(ulong address, byte[] byteArray);

        public abstract bool ReadBytes(ulong address, byte[] array, int count);



        #endregion
    }
    /// <summary>
    /// 远程内存操作的 默认实现 
    /// </summary>
    class RMemoryDefault : MemoryAssistBase, IMemoryAssist
    {

        public override string Description() => "标准内存读写";


        public RMemoryDefault()
        {
            this.ModuleCache = new TtlCache<Int32, List<NormalizedModule>>(TimeSpan.FromSeconds(10.0));

            // Subscribe to process events


            this.ExternalProcess = PAFactory.Default.OpenedProcess;
            PAFactory.Default.Subscribe(this);
        }


        #region 内存分配
        public override ulong AllocateMemory(int size)
        {
            throw new NotImplementedException();
        }

        public override ulong AllocateMemory(int size, ulong allocAddress)
        {
            throw new NotImplementedException();
        }

        public override void FreeMemory(ulong address)
        {
            throw new NotImplementedException();
        }



        #endregion



        #region 内存读写 
        public override bool ReadBytes(ulong address, byte[] array, int count)
        {
            // Allocate the buffer

            if (this.ExternalProcess == null)
            {
                return false;
            }

            bool ok;
            int bytesRead;
            try
            {
                // Read the data from the target process
                ok = Native.NativeMethods.ReadProcessMemory(this.ExternalProcess.Handle, address.ToIntPtr(), array, count, out bytesRead) && count == bytesRead;
            }
            catch
            {
                return false;
            }

            return ok;
        }
        public override byte[] ReadBytes(ulong address, int size, out bool success)
        {
            // Allocate the buffer
            Byte[] buffer = new Byte[size];
            int bytesRead;

            if (this.ExternalProcess == null)
            {
                success = false;
                return null;
            }


            try
            {

                // Read the data from the target process
                success = Native.NativeMethods.ReadProcessMemory(this.ExternalProcess.Handle, address.ToIntPtr(), buffer, size, out bytesRead) && size == bytesRead;
            }
            catch
            {
                success = false;
                return null;
            }

            return buffer;
        }

        public override bool WriteBytes(ulong address, byte[] byteArray)
        {
            MemoryProtectionFlags oldProtection;
            int bytesWritten;

            try
            {
                Native.NativeMethods.VirtualProtectEx(this.ExternalProcess.Handle, address.ToIntPtr(), byteArray.Length, MemoryProtectionFlags.ExecuteReadWrite, out oldProtection);

                // Write the data to the target process
                if (Native.NativeMethods.WriteProcessMemory(this.ExternalProcess.Handle, address.ToIntPtr(), byteArray, byteArray.Length, out bytesWritten))
                {
                    // Check whether all bytes were written
                    if (bytesWritten == byteArray.Length)
                    {
                        return true;
                    }
                }

                Native.NativeMethods.VirtualProtectEx(this.ExternalProcess.Handle, address.ToIntPtr(), byteArray.Length, oldProtection, out oldProtection);
            }
            catch
            {

            }

            return false;
        }


        #endregion
    }
}
